# FMCW雷达信号处理程序
from typing import Tuple
import numpy as np
import matplotlib.pyplot as plt
from apps.fmcw.conf.app_config import AppConfig as AF

class FmcwRsp(object):
    def __init__(self):
        self.name = 'apps.fmcw.core.fmcw_rsp.FmcwRsp'

    @staticmethod
    def range_doppler(RDC:np.ndarray) -> np.ndarray:
        # Perform 2-D FFT on each CPI frame
        RD_frame = RDC[:, 0:AF.numChirps, :] # CPI一帧的数据
        fft2_result = None
        for ridx in range(AF.numTX*AF.numRX):
            fri = np.fft.fft2(RD_frame[:,:,ridx])
            if fft2_result is None:
                fft2_result = np.array([fri]).transpose(1,2,0)
            else:
                fft2_result = np.concatenate((fft2_result, np.array([fri]).transpose(1,2,0)), axis=-1)
        fftshifted_result = None
        for ridx in range(AF.numTX*AF.numRX):
            fsri = np.fft.fftshift(fft2_result[:,:,ridx], axes=(-1,))
            if fftshifted_result is None:
                fftshifted_result = np.array([fsri]).transpose(1,2,0)
            else:
                fftshifted_result = np.concatenate((fftshifted_result, np.array([fsri]).transpose(1,2,0)), axis=-1)
        AF.RDMs[:, :, :] = fftshifted_result
        return AF.RDMs

    @staticmethod
    def draw_range_doppler() -> None:
        # Plotting
        plt.figure()
        # Adjust imshow to match the shape of RDMs
        img_data = 20 * np.log10(np.abs(AF.RDMs[:, :, 0]) / np.max(np.max(np.abs(AF.RDMs[:, :, 0]))))
        img_data = img_data[::-1] # 将图像行进行倒排，第1行变为最后一行，依此类推
        img0 = plt.imshow(img_data, extent=[AF.V[0], AF.V[-1], AF.R[0], AF.R[-1]], aspect='auto')
        # 获取当前图像的颜色限制
        clim = img0.get_clim()
        # 设置颜色轴的限制
        plt.clim(clim[0]/2, 0)
        plt.xlabel('Velocity (m/s)')
        plt.ylabel('Range (m)')
        plt.show()

    @staticmethod
    def ca_cfar(RDMs:np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        RDM_dB = 10*np.log10(abs(RDMs[:,:,1])/np.max(abs(RDMs[:,:,1])))
        numTrain2D = AF.numTrain * AF.numTrain - AF.numGuard * AF.numGuard
        RDM_mask = np.zeros(RDM_dB.shape)
        for r in range(AF.numTrain + AF.numGuard, RDM_mask.shape[0] - (AF.numTrain + AF.numGuard)):
            for d in range(AF.numTrain + AF.numGuard, RDM_mask.shape[1] - (AF.numTrain + AF.numGuard)):                
                Pn = (
                    np.sum(
                        RDM_dB[r-(AF.numTrain+AF.numGuard):r+(AF.numTrain+AF.numGuard), d-(AF.numTrain+AF.numGuard):d+(AF.numTrain+AF.numGuard)]
                    ) - 
                    np.sum(RDM_dB[r-AF.numGuard:r+AF.numGuard, d-AF.numGuard:d+AF.numGuard])
                ) / numTrain2D
                a = numTrain2D * (AF.P_fa**(-1/numTrain2D) - 1)
                threshold = a * Pn
                if RDM_dB[r, d] > threshold and RDM_dB[r, d] > AF.SNR_OFFSET:
                    RDM_mask[r, d] = 1
        cfar_ranges, cfar_dopps = np.where(RDM_mask == 1)  # cfar detected range bins
        # remaining part is for target location estimation
        rem_range = np.zeros(len(cfar_ranges), dtype=int)
        rem_dopp = np.zeros(len(cfar_dopps), dtype=int)
        for i in range(1, len(cfar_ranges)):
            if abs(cfar_ranges[i] - cfar_ranges[i-1]) <= 5 and abs(cfar_dopps[i] - cfar_dopps[i-1]) <= 5:
                rem_range[i] = i  # redundant range indices to be removed
                rem_dopp[i] = i  # redundant doppler indices to be removed
        rem_range = rem_range[rem_range != 0]  # filter zeros
        rem_dopp = rem_dopp[rem_dopp != 0]  # filter zeros
        cfar_ranges = np.delete(cfar_ranges, rem_range)
        cfar_dopps = np.delete(cfar_dopps, rem_dopp)
        K = len(cfar_dopps)  # # of detected targets
        return RDM_mask, cfar_ranges, cfar_dopps, K
    
    @staticmethod
    def draw_ca_cfar(RDM_mask:np.ndarray) -> None:    
        # 绘制CAFR图像 imagesc(V,R,20*log10(abs(RDMs(:,:,1,1))/max(max(abs(RDMs(:,:,1,1))))));
        plt.figure()
        img_data = RDM_mask[::-1]
        img1 = plt.imshow(img_data, extent=[AF.V[0], AF.V[-1], AF.R[0], AF.R[-1]], aspect='auto')
        plt.xlabel('Velocity (m/s)')
        plt.ylabel('Range (m)')
        plt.title('CA-CFAR')
        plt.show()

    @staticmethod
    def music_angle(RDMs:np.ndarray, cfar_ranges:np.ndarray, cfar_dopps:np.ndarray, K:np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        # Angle Estimation - MUSIC Pseudo Spectrum
        # 初始化a1数组，假设其形状为[numTX*numRX, len(ang_ax)]
        a1 = np.zeros((AF.numTX * AF.numRX, len(AF.ang_ax)), dtype=np.complex_)
        # 计算a1
        for k in range(len(AF.ang_ax)):
            # 创建一个从0到numTX*numRX-1的列向量
            idx = np.arange(AF.numTX * AF.numRX)
            # 计算复指数并赋值给a1的第k列
            a1[:, k] = np.exp(-1j * 2 * np.pi * AF.d * (idx * np.sin(np.radians(AF.ang_ax[k]))))
        music_spectrum = np.zeros((K, len(AF.ang_ax)))
        for i in range(K):
            A = RDMs[cfar_ranges[i], cfar_dopps[i], :]
            A = np.expand_dims(A, axis=1)
            Rxx = np.dot(A, A.conj().T)
            D, Q = np.linalg.eig(Rxx)  # Q: eigenvectors (columns), D: eigenvalues
            I = np.argsort(D.real)[::-1]  # Sort eigenvalues in descending order
            Q = Q[:, I]  # Sort the eigenvectors accordingly
            Qs = Q[:, 0]  # Get the signal eigenvector
            Qn = Q[:, 1:]  # Get the noise eigenvectors
            for k in range(len(AF.ang_ax)):
                a1k = np.expand_dims(a1[:, k], 1)
                music_spectrum[i, k] = np.real(np.dot(a1k.conj().T, a1k) / np.dot(np.dot(a1k.conj().T, np.dot(Qn, Qn.conj().T)), a1k))
        return music_spectrum, a1
    
    @staticmethod
    def draw_music_angle(music_spectrum:np.ndarray, K:np.ndarray) -> None:
        # Plotting the MUSIC spectrum
        plt.figure()
        plt.grid(True)
        plt.title('MUSIC Spectrum')
        plt.xlabel('Angle in degrees')
        for k in range(K):
            plt.plot(AF.ang_ax, np.log10(np.abs(music_spectrum[k, :])))
        plt.show()    

    @staticmethod
    def range_angle(RDC:np.ndarray, a1:np.ndarray) -> np.ndarray:
        # 使用 NumPy 的 FFT 函数代替 MATLAB 的 fft
        rangeFFT = np.fft.fft(RDC, axis=0)
        range_az_music = np.zeros((AF.N_range, len(AF.ang_ax)))

        for i in range(AF.N_range):
            # Rxx = np.zeros((AF.numTX * AF.numRX, AF.numTX * AF.numRX), dtype=np.complex_)
            # MATLAB 中的 squeeze 和 sum 函数对应于 NumPy 的 squeeze 和 sum 函数
            A = np.sum(rangeFFT[i, 0:AF.numChirps, :], axis=0).reshape(-1, 1)
            Rxx = np.dot(A, A.conj().T)
            D, Q = np.linalg.eig(Rxx)  # Q: eigenvectors (columns), D: eigenvalues
            I = np.argsort(D.real)[::-1]  # Sort eigenvalues in descending order
            Q = Q[:, I]  # Sort the eigenvectors accordingly
            Qs = Q[:, 0]  # Get the signal eigenvector
            Qn = Q[:, 1:]  # Get the noise eigenvectors
            music_spectrum2 = np.zeros(len(AF.ang_ax))
            for k in range(len(AF.ang_ax)):
                # NumPy 中的矩阵乘法使用 np.dot
                music_spectrum2[k] = np.real(np.dot(a1[:, k].conj().T, a1[:, k]) / np.dot(np.dot(a1[:, k].conj().T, np.dot(Qn, Qn.conj().T)), a1[:, k]))
            range_az_music[i, :] = music_spectrum2
        return range_az_music
    
    @staticmethod
    def draw_range_angle(range_az_music:np.ndarray) -> None:
        plt.figure()
        img_data = 20 * np.log10(np.abs(range_az_music) / np.max(np.abs(range_az_music)))
        img_data = img_data[::-1]
        plt.imshow(img_data, 
                extent=[AF.ang_ax[0], AF.ang_ax[-1], AF.R[0], AF.R[-1]], 
                aspect='auto', 
                cmap='jet')
        plt.colorbar()
        plt.xlabel('Azimuth')
        plt.ylabel('Range (m)')
        plt.title('MUSIC Range-Angle Map')
        plt.clim(vmin=np.min(20 * np.log10(np.abs(range_az_music) / np.max(np.abs(range_az_music)))), 
                vmax=np.max(20 * np.log10(np.abs(range_az_music) / np.max(np.abs(range_az_music)))))
        plt.show()

    @staticmethod
    def point_cloud(music_spectrum:np.ndarray, cfar_ranges:np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        # 找到music_spectrum每行的最大值索引
        I1 = np.argmax(music_spectrum[1, :])
        angle1 = AF.ang_ax[I1]
        I2 = np.argmax(music_spectrum[0, :])
        angle2 = AF.ang_ax[I2]
        # 计算坐标
        coor1 = np.array([cfar_ranges[1] * np.cos(np.deg2rad(angle1)), 
                        cfar_ranges[1] * np.sin(np.deg2rad(angle1)), 0])
        coor2 = np.array([cfar_ranges[0] * np.cos(np.deg2rad(angle2)), 
                        cfar_ranges[0] * np.sin(np.deg2rad(angle2)), 0])
        return coor1, coor2
    
    @staticmethod
    def draw_point_cloud(coor1:np.ndarray, coor2:np.ndarray) -> None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.scatter(coor1[0], coor1[1], coor1[2], s=10, c='m', marker='o', edgecolors='k', linewidth=3)
        ax.scatter(coor2[0], coor2[1], coor2[2], s=10, c='b', marker='o', edgecolors='k', linewidth=3)
        # 设置图形标题和坐标轴标签
        ax.set_title('3D Coordinates (Point Cloud) of the targets')
        ax.set_xlabel('Range (m) X')
        ax.set_ylabel('Range (m) Y')
        ax.set_zlabel('Range (m) Z')
        plt.show()
        print(f'目标１： {coor1[0]}, {coor1[1]}, {coor1[2]}')
        print(f'目标１： {coor2[0]}, {coor2[1]}, {coor2[2]}')
        # 绘制２维散点图
        plt.figure()
        plt.scatter(coor1[0], coor1[1], s=10, c='m', marker='o', edgecolors='k', linewidth=3)
        plt.scatter(coor2[0], coor2[1], s=10, c='b', marker='o', edgecolors='k', linewidth=3)
        # 设置图形标题和坐标轴标签
        plt.title('2D Coordinates (Point Cloud) of the targets')
        plt.xlabel('Range (m) X')
        plt.ylabel('Range (m) Y')
        plt.show()